// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/ipc/dictionary.h>

#include <algorithm>
#include <cstdint>
#include <memory>
#include <set>
#include <unordered_map>
#include <utility>
#include <vector>

#include <nebula/core/array.h>
#include <nebula/array/concatenate.h>
#include <nebula/array/validate.h>
#include <nebula/core/extension_type.h>
#include <nebula/core/record_batch.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <turbo/base/checked_cast.h>
#include <turbo/log/logging.h>

namespace nebula {


    namespace ipc {

        using internal::FieldPosition;

        // ----------------------------------------------------------------------
        // DictionaryFieldMapper implementation

        struct DictionaryFieldMapper::Impl {
            using FieldPathMap = std::unordered_map<FieldPath, int64_t, FieldPath::Hash>;

            FieldPathMap field_path_to_id;

            void ImportSchema(const Schema &schema) {
                ImportFields(FieldPosition(), schema.fields());
            }

            turbo::Status AddSchemaFields(const Schema &schema) {
                if (!field_path_to_id.empty()) {
                    return turbo::invalid_argument_error("Non-empty DictionaryFieldMapper");
                }
                ImportSchema(schema);
                return turbo::OkStatus();
            }

            turbo::Status add_field(int64_t id, std::vector<int> field_path) {
                const auto pair = field_path_to_id.emplace(FieldPath(std::move(field_path)), id);
                if (!pair.second) {
                    return turbo::already_exists_error("Field already mapped to id");
                }
                return turbo::OkStatus();
            }

            turbo::Result<int64_t> GetFieldId(std::vector<int> field_path) const {
                const auto it = field_path_to_id.find(FieldPath(std::move(field_path)));
                if (it == field_path_to_id.end()) {
                    return turbo::not_found_error("Dictionary field not found");
                }
                return it->second;
            }

            int num_fields() const { return static_cast<int>(field_path_to_id.size()); }

            int num_dicts() const {
                std::set<int64_t> uniqueIds;

                for (auto &kv: field_path_to_id) {
                    uniqueIds.insert(kv.second);
                }

                return static_cast<int>(uniqueIds.size());
            }

        private:
            void ImportFields(const FieldPosition &pos,
                              const std::vector<std::shared_ptr<Field>> &fields) {
                for (int i = 0; i < static_cast<int>(fields.size()); ++i) {
                    ImportField(pos.child(i), *fields[i]);
                }
            }

            void ImportField(const FieldPosition &pos, const Field &field) {
                const DataType *type = field.type().get();
                if (type->id() == Type::EXTENSION) {
                    type = turbo::checked_cast<const ExtensionType &>(*type).storage_type().get();
                }
                if (type->id() == Type::DICTIONARY) {
                    InsertPath(pos);
                    // Import nested dictionaries
                    ImportFields(pos,
                                 turbo::checked_cast<const DictionaryType &>(*type).get_value_type()->fields());
                } else {
                    ImportFields(pos, type->fields());
                }
            }

            void InsertPath(const FieldPosition &pos) {
                const int64_t id = field_path_to_id.size();
                const auto pair = field_path_to_id.emplace(FieldPath(pos.path()), id);
                        DKCHECK(pair.second);  // was inserted
                TURBO_UNUSED(pair);
            }
        };

        DictionaryFieldMapper::DictionaryFieldMapper() : impl_(new Impl) {}

        DictionaryFieldMapper::DictionaryFieldMapper(const Schema &schema) : impl_(new Impl) {
            impl_->ImportSchema(schema);
        }

        DictionaryFieldMapper::~DictionaryFieldMapper() {}

        turbo::Status DictionaryFieldMapper::AddSchemaFields(const Schema &schema) {
            return impl_->AddSchemaFields(schema);
        }

        turbo::Status DictionaryFieldMapper::add_field(int64_t id, std::vector<int> field_path) {
            return impl_->add_field(id, std::move(field_path));
        }

        turbo::Result<int64_t> DictionaryFieldMapper::GetFieldId(std::vector<int> field_path) const {
            return impl_->GetFieldId(std::move(field_path));
        }

        int DictionaryFieldMapper::num_fields() const { return impl_->num_fields(); }

        int DictionaryFieldMapper::num_dicts() const { return impl_->num_dicts(); }

        // ----------------------------------------------------------------------
        // DictionaryMemo implementation

        namespace {

            bool HasUnresolvedNestedDict(const ArrayData &data) {
                if (data.type->id() == Type::DICTIONARY) {
                    if (data.dictionary == nullptr) {
                        return true;
                    }
                    if (HasUnresolvedNestedDict(*data.dictionary)) {
                        return true;
                    }
                }
                for (const auto &child: data.child_data) {
                    if (HasUnresolvedNestedDict(*child)) {
                        return true;
                    }
                }
                return false;
            }

        }  // namespace

        struct DictionaryMemo::Impl {
            // Map of dictionary id to dictionary array(s) (several in case of deltas)
            std::unordered_map<int64_t, ArrayDataVector> id_to_dictionary_;
            std::unordered_map<int64_t, std::shared_ptr<DataType>> id_to_type_;
            DictionaryFieldMapper mapper_;

            turbo::Result<decltype(id_to_dictionary_)::iterator> FindDictionary(int64_t id) {
                auto it = id_to_dictionary_.find(id);
                if (it == id_to_dictionary_.end()) {
                    return turbo::not_found_error("Dictionary with id ", id, " not found");
                }
                return it;
            }

            turbo::Result<std::shared_ptr<ArrayData>> ReifyDictionary(int64_t id, MemoryPool *pool) {
                TURBO_MOVE_OR_RAISE(auto it, FindDictionary(id));
                ArrayDataVector *data_vector = &it->second;

                        DKCHECK(!data_vector->empty());
                if (data_vector->size() > 1) {
                    // There are deltas, we need to concatenate them to the first dictionary.
                    ArrayVector to_combine;
                    to_combine.reserve(data_vector->size());
                    // IMPORTANT: At this point, the dictionary data may be untrusted.
                    // We need to validate it, as concatenation can crash on invalid or
                    // corrupted data.  Full validation is necessary for certain types
                    // (for example nested dictionaries).
                    for (const auto &data: *data_vector) {
                        if (HasUnresolvedNestedDict(*data)) {
                            return turbo::unimplemented_error(
                                    "Encountered delta dictionary with an unresolved nested dictionary");
                        }
                        TURBO_RETURN_NOT_OK(::nebula::internal::ValidateArrayFull(*data));
                        to_combine.push_back(make_array(data));
                    }
                    TURBO_MOVE_OR_RAISE(auto combined_dict, concatenate(to_combine, pool));
                    *data_vector = {combined_dict->data()};
                }

                return data_vector->back();
            }
        };

        DictionaryMemo::DictionaryMemo() : impl_(new Impl()) {}

        DictionaryMemo::~DictionaryMemo() {}

        DictionaryFieldMapper &DictionaryMemo::fields() { return impl_->mapper_; }

        const DictionaryFieldMapper &DictionaryMemo::fields() const { return impl_->mapper_; }

        turbo::Result<std::shared_ptr<DataType>> DictionaryMemo::GetDictionaryType(int64_t id) const {
            const auto it = impl_->id_to_type_.find(id);
            if (it == impl_->id_to_type_.end()) {
                return turbo::not_found_error("No record of dictionary type with id ", id);
            }
            return it->second;
        }

        // Returns turbo::not_found_error if dictionary not found
        turbo::Result<std::shared_ptr<ArrayData>> DictionaryMemo::GetDictionary(int64_t id,
                                                                         MemoryPool *pool) const {
            return impl_->ReifyDictionary(id, pool);
        }

        turbo::Status DictionaryMemo::AddDictionaryType(int64_t id,
                                                 const std::shared_ptr<DataType> &type) {
            // AddDictionaryType expects the dict value type
                    DKCHECK_NE(type->id(), Type::DICTIONARY);
            const auto pair = impl_->id_to_type_.emplace(id, type);
            if (!pair.second && !pair.first->second->equals(*type)) {
                return turbo::already_exists_error("Conflicting dictionary types for id ", id);
            }
            return turbo::OkStatus();
        }

        bool DictionaryMemo::HasDictionary(int64_t id) const {
            const auto it = impl_->id_to_dictionary_.find(id);
            return it != impl_->id_to_dictionary_.end();
        }

        turbo::Status DictionaryMemo::AddDictionary(int64_t id,
                                             const std::shared_ptr<ArrayData> &dictionary) {
            const auto pair = impl_->id_to_dictionary_.emplace(id, ArrayDataVector{dictionary});
            if (!pair.second) {
                return turbo::already_exists_error("Dictionary with id ", id, " already exists");
            }
            return turbo::OkStatus();
        }

        turbo::Status DictionaryMemo::AddDictionaryDelta(int64_t id,
                                                  const std::shared_ptr<ArrayData> &dictionary) {
            TURBO_MOVE_OR_RAISE(auto it, impl_->FindDictionary(id));
            it->second.push_back(dictionary);
            return turbo::OkStatus();
        }

        turbo::Result<bool> DictionaryMemo::AddOrReplaceDictionary(
                int64_t id, const std::shared_ptr<ArrayData> &dictionary) {
            ArrayDataVector value{dictionary};

            auto pair = impl_->id_to_dictionary_.emplace(id, value);
            if (pair.second) {
                // Inserted
                return true;
            } else {
                // Update existing value
                pair.first->second = std::move(value);
                return false;
            }
        }

        // ----------------------------------------------------------------------
        // CollectDictionaries implementation

        namespace {

            struct DictionaryCollector {
                const DictionaryFieldMapper &mapper_;
                DictionaryVector dictionaries_;

                turbo::Status WalkChildren(const FieldPosition &position, const DataType &type,
                                    const Array &array) {
                    for (int i = 0; i < type.num_fields(); ++i) {
                        auto boxed_child = make_array(array.data()->child_data[i]);
                        TURBO_RETURN_NOT_OK(Visit(position.child(i), type.field(i), boxed_child.get()));
                    }
                    return turbo::OkStatus();
                }

                turbo::Status Visit(const FieldPosition &position, const std::shared_ptr<Field> &field,
                             const Array *array) {
                    const DataType *type = array->type().get();

                    if (type->id() == Type::EXTENSION) {
                        type = turbo::checked_cast<const ExtensionType &>(*type).storage_type().get();
                        array = turbo::checked_cast<const ExtensionArray &>(*array).storage().get();
                    }
                    if (type->id() == Type::DICTIONARY) {
                        const auto &dict_array = turbo::checked_cast<const DictionaryArray &>(*array);
                        auto dictionary = dict_array.dictionary();

                        // Traverse the dictionary to first gather any nested dictionaries
                        // (so that they appear in the output before their parent)
                        const auto &dict_type = turbo::checked_cast<const DictionaryType &>(*type);
                        TURBO_RETURN_NOT_OK(WalkChildren(position, *dict_type.get_value_type(), *dictionary));

                        // Then record the dictionary itself
                        TURBO_MOVE_OR_RAISE(int64_t id, mapper_.GetFieldId(position.path()));
                        dictionaries_.emplace_back(id, dictionary);
                    } else {
                        TURBO_RETURN_NOT_OK(WalkChildren(position, *type, *array));
                    }
                    return turbo::OkStatus();
                }

                turbo::Status Collect(const RecordBatch &batch) {
                    FieldPosition position;
                    const Schema &schema = *batch.schema();
                    dictionaries_.reserve(mapper_.num_fields());

                    for (int i = 0; i < schema.num_fields(); ++i) {
                        TURBO_RETURN_NOT_OK(Visit(position.child(i), schema.field(i), batch.column(i).get()));
                    }
                    return turbo::OkStatus();
                }
            };

            struct DictionaryResolver {
                const DictionaryMemo &memo_;
                MemoryPool *pool_;

                turbo::Status VisitChildren(const ArrayDataVector &data_vector, FieldPosition parent_pos) {
                    int i = 0;
                    for (const auto &data: data_vector) {
                        // Some data entries may be missing if reading only a subset of the schema
                        if (data != nullptr) {
                            TURBO_RETURN_NOT_OK(VisitField(parent_pos.child(i), data.get()));
                        }
                        ++i;
                    }
                    return turbo::OkStatus();
                }

                turbo::Status VisitField(FieldPosition field_pos, ArrayData *data) {
                    const DataType *type = data->type.get();
                    if (type->id() == Type::EXTENSION) {
                        type = turbo::checked_cast<const ExtensionType &>(*type).storage_type().get();
                    }
                    if (type->id() == Type::DICTIONARY) {
                        TURBO_MOVE_OR_RAISE(const int64_t id,
                                               memo_.fields().GetFieldId(field_pos.path()));
                        TURBO_MOVE_OR_RAISE(data->dictionary, memo_.GetDictionary(id, pool_));
                        // resolve nested dictionary data
                        TURBO_RETURN_NOT_OK(VisitField(field_pos, data->dictionary.get()));
                    }
                    // resolve child data
                    return VisitChildren(data->child_data, field_pos);
                }
            };

        }  // namespace

        turbo::Result<DictionaryVector> CollectDictionaries(const RecordBatch &batch,
                                                     const DictionaryFieldMapper &mapper) {
            DictionaryCollector collector{mapper, {}};
            TURBO_RETURN_NOT_OK(collector.Collect(batch));
            return std::move(collector.dictionaries_);
        }

        namespace internal {

            turbo::Status CollectDictionaries(const RecordBatch &batch, DictionaryMemo *memo) {
                TURBO_RETURN_NOT_OK(memo->fields().AddSchemaFields(*batch.schema()));
                TURBO_MOVE_OR_RAISE(const auto dictionaries,
                                       CollectDictionaries(batch, memo->fields()));
                for (const auto &pair: dictionaries) {
                    TURBO_RETURN_NOT_OK(memo->AddDictionary(pair.first, pair.second->data()));
                }
                return turbo::OkStatus();
            }

        }  // namespace internal

        turbo::Status ResolveDictionaries(const ArrayDataVector &columns, const DictionaryMemo &memo,
                                   MemoryPool *pool) {
            DictionaryResolver resolver{memo, pool};
            return resolver.VisitChildren(columns, FieldPosition());
        }

    }  // namespace ipc
}  // namespace nebula
